#pragma once

#include <iostream>
#include <cassert>
#include <initializer_list>

/**
 * @file RBT.h
 * @author xingyifan
 * @date 2022-12-26 00:17
 *
 * @description: 这部分内联函数用于绘制树
 */

// 重复输出 str n 次
inline void repeat(int n, const char *Str)
{
    while (n > 0)
    {
        std::cout << Str;
        n--;
    }
}

// 获得整数的位数
inline int digit(int n)
{
    if (n < 10 && n >= 0)
        return 1;
    else if (n > -10 && n < 0)
        return 2;
    else
        return digit(n / 10) + 1;
}

/**
 * @file RBT.h
 * @author xingyifan
 * @date 2022-12-26 00:20
 *
 * @description: 类型简化
 */

typedef int Color;
#define RBT_RED 0
#define RBT_BLACK 1

/**
 * @file RBT.h
 * @author xingyifan
 * @date 2022-12-25 22:31
 *
 * @description: 红黑树节点
 */

template <class T>
class RBT_Node
{
public:
    int size;         // 子节点的规模
    T value;          // 节点值
    Color color;      // 节点颜色
    RBT_Node *lc;     // 左孩子
    RBT_Node *rc;     // 右孩子
    RBT_Node *parent; // 父节点

public:
    // 默认构造
    RBT_Node() : size(0), color(RBT_BLACK), value(T()), lc(nullptr), rc(nullptr), parent(nullptr) {}
};

/**
 * @file RBT.h
 * @author xingyifan
 * @date 2022-12-25 22:31
 *
 * @description: 红黑树
 */

template <class T>
class RBT
{
protected:
    RBT_Node<T> *root; // 根节点
    RBT_Node<T> *nil;  // nil 节点作为叶节点方便操作，根节点的父节点是 nil

public:
    // 默认构造和初始化列表构造
    RBT() : nil(new RBT_Node<T>), root(nullptr) { root = nil; }
    RBT(std::initializer_list<T> l) : RBT()
    {
        for (auto it = l.begin(); it != l.end(); ++it)
            insert(*it);
    }
    // 后序遍历删除整棵树
    ~RBT()
    {
        postorderDelete(root);
        delete nil;
    }

    // 获得树的规模
    int size() const { return root->size; }

    /**
     * @file RBT.h
     * @author xingyifan
     * @date 2022-12-25 22:46
     *
     * @description: 节点插入和删除算法
     */

    void insert(T value)
    {
        // 此时树是空的，要在根节点处创建一个新的
        if (root == nil)
        {
            root = new RBT_Node<T>;
            root->value = value;
            root->parent = nil;
            root->lc = nil;
            root->rc = nil;
            root->size = 1;
            root->color = RBT_BLACK;
            return;
        }

        // 否则从根节点出发向下，利用比较确定向左还是向右，沿途的节点 size 都加 1
        RBT_Node<T> *p = root;
        RBT_Node<T> *pos = nil; // 记录最终插入位置
        while (true)
        {
            p->size++;
            // 向左，遇到叶节点就创建新节点，然后退出循环
            if (value < p->value)
            {
                if (p->lc != nil)
                    p = p->lc;
                else
                {
                    p->lc = new RBT_Node<T>;
                    pos = p->lc;
                    break;
                }
            }
            // 向右，遇到叶节点就创建新节点，然后退出循环
            else
            {
                if (p->rc != nil)
                    p = p->rc;
                else
                {
                    p->rc = new RBT_Node<T>;
                    pos = p->rc;
                    break;
                }
            }
        }

        // 插入点赋值，设为红节点，不改变黑高
        pos->value = value;
        pos->parent = p;
        pos->rc = nil;
        pos->lc = nil;
        pos->size = 1;
        pos->color = RBT_RED;

        // 调整插入节点
        RBT_Node<T> *me = pos;
        RBT_Node<T> *father = me->parent;

        // 只有当父节点是红节点才需要调整
        while (father->color == RBT_RED)
        {
            // 父节点是左节点的情况
            if (father->parent->lc == father)
            {
                RBT_Node<T> *uncle = father->parent->rc;
                // case 1：父节点和叔节点都是红色，说明父节点的父节点是黑节点
                if (uncle->color == RBT_RED)
                {
                    father->color = RBT_BLACK;
                    uncle->color = RBT_BLACK;
                    me = father->parent;
                    father = me->parent;
                    me->color = RBT_RED;
                }
                // case 2：叔节点是黑色
                else
                {
                    if (father->lc != me)
                    {
                        leftRotate(father);
                        me = me->lc;
                        father = me->parent;
                    }
                    father->color = RBT_BLACK;
                    father->parent->color = RBT_RED;
                    rightRotate(father->parent);
                }
            }
            else
            {
                RBT_Node<T> *uncle = father->parent->lc;
                // case 1：父节点和叔节点都是红色，说明父节点的父节点是黑节点
                if (uncle->color == RBT_RED)
                {
                    father->color = RBT_BLACK;
                    uncle->color = RBT_BLACK;
                    me = father->parent;
                    father = me->parent;
                    me->color = RBT_RED;
                }
                // case 2：叔节点是黑色
                else
                {
                    if (father->rc != me)
                    {
                        rightRotate(father);
                        me = me->rc;
                        father = me->parent;
                    }
                    father->color = RBT_BLACK;
                    father->parent->color = RBT_RED;
                    leftRotate(father->parent);
                }
            }
        }
        // 根节点一定是黑色
        root->color = RBT_BLACK;
    }
    bool remove(T value)
    {
        // 先尝试找到 value 对应的节点，如果找不到就说明不存在
        RBT_Node<T> *p = findNode(value);
        if (p == nil)
            return false;

        // 检查左右节点
        RBT_Node<T> *lp = p->lc;
        RBT_Node<T> *rp = p->rc;
        RBT_Node<T> *adj = nil;
        Color origin = p->color;

        // size调整
        pathReduce(root, p);
        if (rp == nil)
        {
            adj = lp;
            // 用 lp 替换 p
            replace(p, lp);

            // p 大小不变，要减去 1
            lp->size = p->size - 1;
            delete p;
        }
        else if (lp == nil)
        {
            adj = rp;
            // 用 rp 替换 p
            replace(p, rp);

            // p 大小不变，要减去 1
            rp->size = p->size - 1;
            delete p;
        }
        else
        {
            // 如果有右节点，就选择右边子树中的最小节点作为新节点
            RBT_Node<T> *pos = minNode(rp);
            origin = pos->color;
            adj = pos->rc;

            // size 调整
            pathReduce(p, pos);

            // 如果新节点不是右节点，就连接新的父节点和洗的右节点
            if (pos != rp)
            {
                replace(pos, pos->rc);
                pos->rc = rp;
                rp->parent = pos;
            }
            // 连接新节点和 p 的父节点和 p 的左右节点
            pos->lc = lp;
            lp->parent = pos;

            replace(p, pos);
            pos->size = p->size;
            pos->color = p->color;
            delete p;
        }

        // 如果删除的节点是黑色，就需要调整节点来平衡黑高
        if (origin == RBT_BLACK)
        {
            // 如果 me 是红色，则变为黑色，然后循环；如果 me 是根，同样
            RBT_Node<T> *me = adj;
            while (me != root && me->color == RBT_BLACK)
            {
                // me 是左节点
                if (me == me->parent->lc)
                {
                    RBT_Node<T> *brother = me->parent->rc;
                    // 如果兄弟是红色，就改成黑色
                    if (brother->color == RBT_RED)
                    {
                        me->parent->color = RBT_RED;
                        brother->color = RBT_BLACK;
                        leftRotate(me->parent);
                        brother = me->parent->rc;
                    }
                    // 如果兄弟的两个孩子都是黑色，就让兄弟是红色，me 向上变为父节点
                    if (brother->lc->color == RBT_BLACK && brother->rc->color == RBT_BLACK)
                    {
                        brother->color = RBT_RED;
                        me = me->parent;
                    }
                    else
                    {
                        // 如果兄弟的右节点是黑色，就修改它，让左节点是黑色
                        if (brother->rc->color == RBT_BLACK)
                        {
                            brother->lc->color = RBT_BLACK;
                            brother->color = RBT_RED;
                            rightRotate(brother);
                            brother = brother->parent;
                        }
                        brother->color = me->parent->color;
                        me->parent->color = RBT_BLACK;
                        brother->rc->color = RBT_BLACK;
                        leftRotate(me->parent);
                        me = root;
                    }
                }
                // 和左节点的情况对称
                else
                {
                    RBT_Node<T> *brother = me->parent->lc;
                    if (brother->color == RBT_RED)
                    {
                        me->parent->color = RBT_RED;
                        brother->color = RBT_BLACK;
                        rightRotate(me->parent);
                        brother = me->parent->lc;
                    }
                    if (brother->rc->color == RBT_BLACK && brother->lc->color == RBT_BLACK)
                    {
                        brother->color = RBT_RED;
                        me = me->parent;
                    }
                    else
                    {
                        if (brother->lc->color == RBT_BLACK)
                        {
                            brother->rc->color = RBT_BLACK;
                            brother->color = RBT_RED;
                            leftRotate(brother);
                            brother = brother->parent;
                        }
                        brother->color = me->parent->color;
                        me->parent->color = RBT_BLACK;
                        brother->lc->color = RBT_BLACK;
                        rightRotate(me->parent);
                        me = root;
                    }
                }
            }
            // 注意最后自己一定是黑色
            me->color = RBT_BLACK;
        }
        return true;
    }

    /**
     * @file RBT.h
     * @author xingyifan
     * @date 2022-12-25 23:18
     *
     * @description: 元素索引方法
     */

    T operator[](int i) const
    {
        // 不允许越界
        assert(i >= 0 && i < root->size);
        RBT_Node<T> *p = root;

        int index = p->lc->size;
        while (index != i)
        {
            if (index < i)
            {
                p = p->rc;
                index += p->lc->size + 1;
            }
            else
            {
                p = p->lc;
                index -= p->rc->size + 1;
            }
        }
        return p->value;
    }
    T &operator[](int i)
    {
        // 不允许越界
        assert(i >= 0 && i < root->size);
        RBT_Node<T> *p = root;

        int index = p->lc->size;
        while (index != i)
        {
            if (index < i)
            {
                p = p->rc;
                index += p->lc->size + 1;
            }
            else
            {
                p = p->lc;
                index -= p->rc->size + 1;
            }
        }
        return p->value;
    }

    // 计算给定值在树中的位置
    int indexOf(T value) const
    {
        if (root == nil)
            return -1;

        RBT_Node<T> *p = root;
        int index = p->lc->size;
        while (p != nil)
        {
            if (p->value < value)
            {
                p = p->rc;
                index += p->lc->size + 1;
            }
            else if (p->value > value)
            {
                p = p->lc;
                index -= p->rc->size + 1;
            }
            else
                return index;
        }
        return -1;
    }

    /**
     * @file RBT.h
     * @author xingyifan
     * @date 2022-12-25 23:17
     *
     * @description: 中序/后序遍历
     */

    void inorderWalk(std::ostream &out, RBT_Node<T> *node) const
    {
        if (node != nil)
        {
            // 先输出左子树，再输出当前节点，最后输出右子树
            inorderWalk(out, node->lc);
            out << node->value << " ";
            inorderWalk(out, node->rc);
        }
    }
    void postorderDelete(RBT_Node<T> *node)
    {
        if (node != nil)
        {
            // 先删除左子树，再删除右子树，最后删除当前节点
            postorderDelete(node->lc);
            postorderDelete(node->rc);
            delete node;
            node = nil;
        }
    }

    /**
     * @file RBT.h
     * @author xingyifan
     * @date 2022-12-26 00:05
     *
     * @description: 顺序输出
     */

    friend std::ostream &operator<<(std::ostream &out, const RBT &bt)
    {
        bt.inorderWalk(out, bt.root);
        return out;
    }

    /**
     * @file RBT.h
     * @author xingyifan
     * @date 2022-12-25 23:26
     *
     * @description: 高度查询
     */

    int height(RBT_Node<T> *node = nullptr) const
    {
        // 默认设为根节点
        if (node == nullptr)
            node = root;
        // 高度是左右子树中最高的高度 +1
        if (node != nil)
            return std::max(height(node->lc), height(node->rc)) + 1;
        return 0;
    }
    int blackHeight(RBT_Node<T> *node = nullptr) const
    {
        // 默认设为根节点
        if (node == nullptr)
            node = root;
        if (node != nil)
        {
            int h = std::max(blackHeight(node->lc), blackHeight(node->rc));
            if (node->color == RBT_BLACK)
                h++;
            return h;
        }
        return 0;
    }

    /**
     * @file RBT.h
     * @author xingyifan
     * @date 2022-12-25 23:36
     *
     * @description: 左旋和右旋
     */

    void leftRotate(RBT_Node<T> *node)
    {
        // 没有右节点就退出
        if (node->rc == nil)
            return;

        RBT_Node<T> *rc = node->rc;
        // 调整 size 大小
        node->size -= rc->rc->size + 1;
        rc->size += node->lc->size + 1;

        rc->parent = node->parent;
        // 如果恰好是根节点，就重置根节点
        if (node == root)
            root = rc;
        else if (node == node->parent->rc)
            node->parent->rc = rc;
        else
            node->parent->lc = rc;

        node->rc = rc->lc;
        rc->lc->parent = node;
        rc->lc = node;
        node->parent = rc;
    }
    void rightRotate(RBT_Node<T> *node)
    {
        // 没有左节点就退出
        if (node->lc == nil)
            return;

        RBT_Node<T> *lc = node->lc;
        // 调整 size 大小
        node->size -= lc->lc->size + 1;
        lc->size += node->rc->size + 1;

        lc->parent = node->parent;
        // 如果恰好是根节点，就重置根节点
        if (node == root)
            root = lc;
        else if (node == node->parent->lc)
            node->parent->lc = lc;
        else
            node->parent->rc = lc;

        node->lc = lc->rc;
        lc->rc->parent = node;
        lc->rc = node;
        node->parent = lc;
    }

    /**
     * @file RBT.h
     * @author xingyifan
     * @date 2022-12-25 22:54
     *
     * @description: 调整节点规模，从 tail 一路向上，直到 head 为止，减少路径上节点的 size
     */

    void pathReduce(RBT_Node<T> *head, RBT_Node<T> *tail)
    {
        while (tail != head)
        {
            tail = tail->parent;
            tail->size--;
        }
    }

    /**
     * @file RBT.h
     * @author xingyifan
     * @date 2022-12-25 22:58
     *
     * @description: 节点替换，用 v 替换 u 的位置
     */

    void replace(RBT_Node<T> *u, RBT_Node<T> *v)
    {
        if (u == root)
            root = v;
        else if (u == u->parent->lc)
            u->parent->lc = v;
        else
            u->parent->rc = v;
        v->parent = u->parent;
    }

    /**
     * @file RBT.h
     * @author xingyifan
     * @date 2022-12-25 22:59
     *
     * @description: 节点查找
     */

    RBT_Node<T> *findNode(T value)
    {
        RBT_Node<T> *p = root;
        while (p != nil)
        {
            if (p->value < value)
                p = p->rc;
            else if (p->value > value)
                p = p->lc;
            else
                break;
        }
        return p;
    }
    RBT_Node<T> *maxNode(RBT_Node<T> *p)
    {
        if (p == nil)
            return nil;
        while (p->rc != nil)
            p = p->rc;
        return p;
    }
    RBT_Node<T> *minNode(RBT_Node<T> *p)
    {
        if (p == nil)
            return nil;
        while (p->lc != nil)
            p = p->lc;
        return p;
    }

    T min() { return minNode(root)->value; }
    T max() { return maxNode(root)->value; }

    // 查找给定节点的后/前一个
    RBT_Node<T> *next(RBT_Node<T> *p)
    {
        // 如果有右节点，就返回右边最小的节点
        if (p->rc != nil)
            return minNode(p->rc);

        // 否则就向上查找，直到 pos==nil 或者 pos 从左边返回
        RBT_Node<T> *pos = p->parent;
        while (pos != nil && pos->lc != p)
        {
            p = pos;
            pos = p->parent;
        }
        return pos;
    }
    RBT_Node<T> *prev(RBT_Node<T> *p)
    {
        // 如果有左节点，就返回左边最大的节点
        if (p->lc != nil)
            return maxNode(p->lc);

        // 否则就向上查找，直到 pos==nil 或者 pos 从右边返回
        RBT_Node<T> *pos = p->parent;
        while (pos != nil && pos->rc != p)
        {
            p = pos;
            pos = p->parent;
        }
        return pos;
    }

    // 查找第 depth 层的 i 号节点
    RBT_Node<T> *floorNode(int depth, int i) const
    {
        RBT_Node<T> *p = root;

        // 获取全局位置
        double left = 0;
        double right = 1 << (depth - 1);
        int floor = 1;
        while (floor != depth && p != nil)
        {
            if (i < (left + right) / 2)
            {
                p = p->lc;
                right = (left + right) / 2;
            }
            else
            {
                p = p->rc;
                left = (left + right) / 2;
            }
            floor++;
        }
        return p;
    }

    /**
     * @file RBT.h
     * @author xingyifan
     * @date 2022-12-26 00:25
     *
     * @description: 红黑树绘制函数
     */

    bool drawTree() const
    {
        // 高度为零，什么都不做
        int h = height();
        if (h == 0)
            return false;

        // 打印第一行
        int blank = (1 << h) - 2;
        repeat(blank, " ");
        std::cout << root->value << std::endl;

        // 每行前面的空格数 2^n - 2 个，从第二行开始
        for (int i = 1; i < h; i++)
        {
            const int num = 1 << i;
            int c = num;

            // 暂存这一层的节点，绘制前置空格
            RBT_Node<T> **record = new RBT_Node<T> *[num];
            blank = (1 << (h - i)) - 2;
            repeat(blank, " ");

            // 输出树枝所在的行
            while (c > 0)
            {
                record[num - c] = floorNode(i + 1, num - c);
                blank = (1 << (h - i)) - 1;

                // 获得此行元素
                RBT_Node<T> *p = record[num - c];
                if (p != nil)
                {
                    // 输出左端点
                    if (c % 2 == 0)
                    {
                        std::cout << "|";
                        repeat(blank + 1, "-");
                    }
                    // 输出右端点
                    else
                    {
                        // 回退一格，覆盖左端点为 nil 时产生的空格
                        std::cout << "\b-";
                        repeat(blank, "-");
                        std::cout << "|";
                        repeat(2 * blank + 1, " ");
                    }
                }
                // 当此节点为空，就用空格填充
                else if (c > 1)
                {
                    // 左端点
                    if (c % 2 == 0)
                        repeat(blank + 2, " ");
                    else
                        repeat(3 * blank + 2, " ");
                }
                c--;
            }
            // 此层输出完毕
            std::cout << std::endl;

            c = num;
            blank = (1 << (h - i)) - 2;
            int d = 0;
            // 输出节点所在的行
            while (c > 0)
            {
                repeat(blank, " ");
                repeat(d, "\b");
                // 只需要在这里插入节点输出
                if (record[num - c] != nil)
                {
                    // 将可能的浮点型直接转为整型，便于格式化
                    int intValue = int(record[num - c]->value);
                    Color color = record[num - c]->color;
                    std::cout << intValue;

                    // 输出节点颜色
                    if (color == RBT_BLACK)
                        std::cout << "b";
                    else
                        std::cout << "r";

                    // 这里要得到该数值的位数
                    d = digit(intValue);
                }
                // 此节点为空，什么都不输出
                else
                {
                    std::cout << " ";
                    d = 0;
                }
                // 相邻间距 2^(n+1) - 1
                blank = (1 << (h - i + 1)) - 1;
                c--;
            }
            // 移除临时记录表
            delete[] record;
            std::cout << std::endl;
        }
        return true;
    }
};